# 三维Laplacian方程有限差分求解
# \laplacian u = b
# 可能有bug
# Gitee Repo

import numpy as np
import scipy

import plotly.graph_objects as go
from plotly.offline import plot

L = 2;
dx = 0.1;

x,y,z = np.meshgrid(np.arange(-L,L,dx),np.arange(-L,L,dx),np.arange(-L,L,dx),indexing='ij');
n = int(x.shape[0]);

# 构造laplacian算符的矩阵
def construct_A(x,y,z,n,L):
    print('Constructing A...')
    
    A = scipy.sparse.dok_matrix((n**3,n**3))
    for i in range(n):
        for j in range(n):
            for k in range(n):
                ind = np.zeros((7,1))
                ind[0] = np.ravel_multi_index((i,j,k),(n,n,n))
                
                # 如果在边界上
                if i == 0 or i == n-1 or j == 0 or j == n-1 or k == 0 or k == n-1:
                    A[ind[0],ind[0]] = 1
                else:
                    ind[1] = np.ravel_multi_index((i+1,j,k),(n,n,n))
                    ind[2] = np.ravel_multi_index((i-1,j,k),(n,n,n))
                    ind[3] = np.ravel_multi_index((i,j+1,k),(n,n,n))
                    ind[4] = np.ravel_multi_index((i,j-1,k),(n,n,n))
                    ind[5] = np.ravel_multi_index((i,j,k+1),(n,n,n))
                    ind[6] = np.ravel_multi_index((i,j,k-1),(n,n,n))
                    
                    A[ind[0],ind[0]]=-6
                    for l in range(1,7):
                        A[ind[0],ind[l]]=1
    return A

# 改善效率
def construct_A2():
    print('Constructing A2...')

    A = scipy.sparse.dok_matrix((n**3,n**3))
    i,j,k = np.meshgrid(np.arange(1,n-1),np.arange(1,n-1),np.arange(1,n-1),indexing='ij')
    i,j,k = i.flatten(),j.flatten(),k.flatten()

    ind = np.zeros(((n-2)**3,7))
    ind[:,0] = np.ravel_multi_index((i,j,k),(n,n,n))
    ind[:,1] = np.ravel_multi_index((i+1,j,k),(n,n,n))
    ind[:,2] = np.ravel_multi_index((i-1,j,k),(n,n,n))
    ind[:,3] = np.ravel_multi_index((i,j+1,k),(n,n,n))
    ind[:,4] = np.ravel_multi_index((i,j-1,k),(n,n,n))
    ind[:,5] = np.ravel_multi_index((i,j,k+1),(n,n,n))
    ind[:,6] = np.ravel_multi_index((i,j,k-1),(n,n,n))
    
    A[ind[:,0],ind[:,0]]=-6
    for l in range(1,7):
        A[ind[:,0],ind[:,l]]=1

    i,j = np.meshgrid(np.arange(0,n),np.arange(0,n),indexing='ij')
    i,j = i.flatten(),j.flatten()
    k0 = np.zeros_like(i)
    k1 = (n-1)*np.ones_like(i)
    ind = np.zeros((n**2,1))
    
    ind[:,0] = np.ravel_multi_index((k0,i,j),(n,n,n))
    A[ind[:,0],ind[:,0]] = 1
    
    ind[:,0] = np.ravel_multi_index((i,k0,j),(n,n,n))
    A[ind[:,0],ind[:,0]] = 1
    
    ind[:,0] = np.ravel_multi_index((i,j,k0),(n,n,n))
    A[ind[:,0],ind[:,0]] = 1
    
    ind[:,0] = np.ravel_multi_index((k1,i,j),(n,n,n))
    A[ind[:,0],ind[:,0]] = 1
    
    ind[:,0] = np.ravel_multi_index((i,k1,j),(n,n,n))
    A[ind[:,0],ind[:,0]] = 1
    
    ind[:,0] = np.ravel_multi_index((i,j,k1),(n,n,n))
    A[ind[:,0],ind[:,0]] = 1

    return A

# 构造边界条件
def construct_b(x,y,z,n,L):
    print('Constructing b...')
    b = np.zeros((n**3,1))
    
    # 边界条件
    #for i in range(n):
    #    for j in range(n):
    #        k = n-1
    #        ind = np.ravel_multi_index((i,j,k),(n,n,n))
    #        b[ind] = np.cos(np.pi/2/L*x[i,j,k])*np.cos(np.pi/2/L*y[i,j,k])
    
    ind = np.ravel_multi_index((int(0.2*n),int(n/2),int(n/2)),(n,n,n))
    b[ind]=-0.5
    return b

# 改善效率
def construct_b2():
    print('Constructing b2...')
    b = np.zeros((n**3,1))

    #i,j = np.meshgrid(np.arange(0,n),np.arange(0,n),indexing='ij')
    #i,j = i.flatten(),j.flatten()
    #k = (n-1)*np.ones_like(i)
    #ind = np.ravel_multi_index((i,j,k),(n,n,n))
    #b[ind,0] = np.cos(np.pi/2/L*np.take(x,ind))*np.cos(np.pi/2/L*np.take(y,ind))
    
    ind = np.ravel_multi_index((int(0.2*n),int(n/2),int(n/2)),(n,n,n))
    b[ind] = -0.5
    return b

A = construct_A2()
b = construct_b2()

print('Coverting to csr matrix...')
A = scipy.sparse.csr_matrix(A)

print('Solving u...')
#_u = scipy.sparse.linalg.spsolve(A,b)
_u = scipy.sparse.linalg.gmres(A,b)[0] #gmres应该更快一点

u = _u.reshape(n,n,n)

scipy.io.savemat('res.mat',{'x':x,'y':y,'z':z,'u':u})

fig = go.Figure(data=go.Volume(
    x=x.flatten(),
    y=y.flatten(),
    z=z.flatten(),
    value=u.flatten(),
    opacity=0.1, # needs to be small to see through all surfaces
    surface_count=20, # needs to be a large number for good volume rendering
    ))
plot(fig)
